#include <optix.h>
#include <optixu/optixu_math_namespace.h>
#include "device_util.h"
#include "random.h" 

using namespace optix; 

rtDeclareVariable(uint2, launch_index, rtLaunchIndex, );
rtBuffer<float4, 2>				result_buffer;
rtBuffer<ViewRayData, 2>  hit_buffer;  
rtBuffer<PhotonRec, 1>		kd_buffer; 
rtBuffer<uint2, 2>        seed_buffer;

rtDeclareVariable(float, flux_div_constant, , );
 
rtDeclareVariable(uint2, launch_dim,   rtLaunchDim, );
  
rtDeclareVariable(Light, light , , );
rtDeclareVariable(rtObject,      top_object, , );

RT_PROGRAM void cast() 
{
	uint2 lsz = launch_dim;
	uint2 li = launch_index;
	ViewRayData hit = hit_buffer[li];
	PhotonRec photon;

	int mdep = 22;
	int stack[22];
	int top = 0, pid;
	stack[top] = 0;
	float r2 = 400.f;
	float d;
	float3 accum_flux = make_float3(0.0f);
	float3 diff;
			
	while ( top >= 0 ) {
		pid = stack[top--]; 
		photon = kd_buffer[pid];
		if ( photon.axis < 0 ) {
			continue;
		}
		diff = photon.pos - hit.position; 
		if (dot(diff, diff) < r2) {
			if (dot(photon.normal, hit.normal) > 0.001f) {
				accum_flux += photon.energy * hit.color;
			}
		}

		switch ( photon.axis ) {				
		case 0:
			d = diff.x;
			break;
		case 1:
			d = diff.y;
			break;
		case 2:
			d = diff.z;
			break;
		} 
				 
		if ( d*d < r2 ) { 
									
			if ( top < mdep && photon.axis != 3 ) {
				stack[++top] = pid * 2 + 1;
				stack[++top] = pid * 2 + 2;
			}
		}	

		else {
			int child = d < 0.0f ? 1 : 0;
			if ( top < mdep && photon.axis != 3 ) {
				stack[++top] = pid * 2 + child + 1;
			}
		}	 
			
	} //while - kd search
			 
 
	float3 indirect_light = 1.0f / ( M_PIf * r2 ) * accum_flux / (flux_div_constant); 
	float3 direct_flux = make_float3(0.f);
	
	uint2  seed   = seed_buffer[launch_index];

	int ndirect = 100;
	for (int i=0; i<ndirect; ++i) {
		float3 point_on_light; 
  
		float2 sample = make_float2( rnd( seed.x ), rnd( seed.y ) ); 
		point_on_light = light.anchor + sample.x*light.v1 + sample.y*light.v2;

		float3 to_light    = point_on_light - hit.position;
		float  light_dist  = length( to_light );
	 
		to_light = normalize(to_light);
		float  n_dot_l     = fmaxf( 0.0f, dot( hit.normal, to_light ) );
		float  light_atten = n_dot_l;

		ShadowRayData prd;
		prd.atten = 1.0f;

		if ( light_atten > 1e-5f ) {  
			optix::Ray shadow_ray( hit.position, to_light, 3, 1e-2, light_dist - 1e-1 );
			rtTrace( top_object, shadow_ray, prd );
			light_atten *= prd.atten * dot( -to_light, light.direction ); 
		} 
		light_atten /= light_dist*light_dist;	
		if( light_atten < 0.0f ) light_atten = 0.0f;  
	
		direct_flux += light.power * light_atten * hit.color;
	}

	direct_flux /= ndirect;
	 
	result_buffer[launch_index] = make_float4(indirect_light, 0.0f); //make_float4(direct_flux * 0.6f + indirect_light + make_float3(0.1f, 0.1f, 0.1f) * hit.color, 0.f);
}

rtDeclareVariable(ShadowRayData, prd, rtPayload, );

RT_PROGRAM void any_hit()
{
	prd.atten = 0.f;
}

